package com.stone.starter.core.script;

import javax.tools.*;
import javax.tools.JavaFileObject.Kind;
import java.io.ByteArrayOutputStream;
import java.io.FilterOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.net.URI;
import java.nio.CharBuffer;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 内存代码管理
 *
 * @author Mr_wenpan@163.com 2021/9/7 2:15 下午
 */
class MemoryJavaCodeManager extends ForwardingJavaFileManager<JavaFileManager> {

    /**
     * 字节码缓存（key为类名，value为类的字节码二进制内容）
     */
    private final Map<String, byte[]> classBytes;

    MemoryJavaCodeManager(JavaFileManager fileManager) {
        super(fileManager);
        classBytes = new ConcurrentHashMap<>();
    }

    /**
     * 获取类的字节码
     */
    Map<String, byte[]> getClassBytes() {
        return new HashMap<>(classBytes);
    }

    @Override
    public void flush() {
    }

    @Override
    public void close() {
        classBytes.clear();
    }

    @Override
    public JavaFileObject getJavaFileForOutput(Location location,
                                               String className,
                                               Kind kind,
                                               FileObject sibling) throws IOException {
        if (kind == Kind.CLASS) {
            return new MemoryOutputJavaFileObject(className);
        } else {
            return super.getJavaFileForOutput(location, className, kind, sibling);
        }
    }

    JavaFileObject makeStringSource(String name, String code) {
        return new MemoryInputJavaFileObject(name, code);
    }

    static class MemoryInputJavaFileObject extends SimpleJavaFileObject {

        final String code;

        MemoryInputJavaFileObject(String name, String code) {
            super(URI.create("string:///" + name), Kind.SOURCE);
            this.code = code;
        }

        @Override
        public CharBuffer getCharContent(boolean ignoreEncodingErrors) {
            return CharBuffer.wrap(code);
        }
    }

    /**
     * 用MemoryOutputJavaFileObject替换JDK默认的SimpleJavaFileObject，
     * 以便在接收到编译器生成的byte[]内容时，不写入class文件，而是直接保存在内存中
     */
    class MemoryOutputJavaFileObject extends SimpleJavaFileObject {

        final String name;

        MemoryOutputJavaFileObject(String name) {
            super(URI.create("string:///" + name), Kind.CLASS);
            this.name = name;
        }

        @Override
        public OutputStream openOutputStream() {
            return new FilterOutputStream(new ByteArrayOutputStream()) {

                @Override
                public void close() throws IOException {
                    out.close();
                    ByteArrayOutputStream bos = (ByteArrayOutputStream) out;
                    classBytes.put(name, bos.toByteArray());
                }
            };
        }

    }
}
